

#pragma once
#include "FrameAllocator.h"

namespace FlameEngine
{
	////////////////////////////////////////////////
	//FrameAllocMemChunk
	FrameAllocMemChunk::FrameAllocMemChunk(size_t size)
		:m_chunkSize(size),
		m_pDataStart(nullptr),
		m_pDataCurrent(nullptr)
	{
		m_pDataStart = new uint8_t[m_chunkSize];
		m_pDataCurrent = m_pDataStart;
		m_freeByteCount = m_chunkSize;
	}

	FrameAllocMemChunk::~FrameAllocMemChunk()
	{
		if(m_pDataStart)
			delete[] m_pDataStart;

		m_pDataStart = nullptr;
		m_pDataCurrent = nullptr;
	}

	void* FrameAllocMemChunk::Allocate(size_t size, size_t align, int32_t& allocByteCount)
	{
		size_t s = m_freeByteCount;
		void* p = m_pDataCurrent;
		void* pResult = std::align(align, size, p, s);
		if(pResult)
		{
			m_pDataCurrent = (uint8_t*)pResult + size;
			allocByteCount = (m_freeByteCount - s) + size;
			m_freeByteCount -= allocByteCount;
		}
		return pResult;
	}

	/////////////////////////////////////////////////
	//FrameAllocator
	FrameAllocator::FrameAllocator()
		:m_chunkSize(2 * 1024 * 1024),
		m_curChunkIdx(-1),
		m_usedByteNum(0)
	{
	}

	FrameAllocator::~FrameAllocator()
	{
		for(size_t i = 0; i < m_chunks.size(); i++)
			delete m_chunks[i];

		m_chunks.clear();
	}

	void* FrameAllocator::Allocate(size_t size, int32_t align)
	{
		if (align < 4)
			align = 4;

		if (m_curChunkIdx < 0)
		{
			m_chunks.push_back(new FrameAllocMemChunk(m_chunkSize));
			m_curChunkIdx = m_chunks.size() - 1;
		}

		int32_t allocSize = 0;
		void* pAlloc = m_chunks[m_curChunkIdx]->Allocate(size, align, allocSize);
		if (pAlloc == nullptr)
		{
			//can not fit in current chunk,create new one
			m_chunks.push_back(new FrameAllocMemChunk(m_chunkSize));
			m_curChunkIdx = m_chunks.size() - 1;
			pAlloc = m_chunks[m_curChunkIdx]->Allocate(size, align, allocSize);
		}

		assert(pAlloc != nullptr && "frame allocator failed");

		m_usedByteNum += allocSize;
		return pAlloc;
	}

	void FrameAllocator::Clear()
	{
		//1mb
		size_t minChunkSize = 1024 * 1024;

		size_t freeSize = m_chunks.size() * m_chunkSize - m_usedByteNum;
		if(freeSize > minChunkSize * 4)
		{
			//too much empty space, release some
			size_t increase = minChunkSize - m_usedByteNum % minChunkSize;
			increase = (increase >= (minChunkSize * 0.3)) ? increase : (increase + minChunkSize);
			m_chunkSize = m_usedByteNum + increase;

			for(size_t i = 0; i < m_chunks.size(); i++)
				delete m_chunks[i];

			m_chunks.clear();
			m_curChunkIdx = -1;
		}
		else if(m_chunks.size() >= 4)
		{
			//too much chunk combined to one
			for(size_t i = 0; i < m_chunks.size(); i++)
				delete m_chunks[i];

			m_chunks.clear();
			m_curChunkIdx = -1;

			size_t totalAllocByte = m_chunks.size() * m_chunkSize;
			size_t delta = totalAllocByte - m_usedByteNum;
			if (delta < (minChunkSize*0.3f))
				totalAllocByte += minChunkSize;
				
			m_chunkSize = totalAllocByte;
		}
		else
		{
			for(size_t i = 0; i < m_chunks.size(); i++)
				m_chunks[i]->Clear();

			if (m_chunks.size() > 0)
				m_curChunkIdx = 0;
		}
		m_usedByteNum = 0;
	}

	/////////////////////////////////////////////////
	//FrameAllocatorThreadSafe
	FrameAllocatorThreadSafe::FrameAllocatorThreadSafe()
	{
	}

	FrameAllocatorThreadSafe::~FrameAllocatorThreadSafe()
	{
		for(auto& alloc : m_threadAlloc)
			delete alloc.second;

		m_threadAlloc.clear();
	}

	FrameAllocator* FrameAllocatorThreadSafe::GetAllocator()
	{
		std::thread::id id = std::this_thread::get_id();
		for(auto& alloc : m_threadAlloc)
		{
			if(alloc.first == id)
				return alloc.second;
		}

		std::lock_guard<std::mutex> lock(m_mutex);
		//double check
		for(auto& alloc : m_threadAlloc)
		{
			if(alloc.first == id)
				return alloc.second;
		}

		FrameAllocator* alloc = new FrameAllocator();
		m_threadAlloc.emplace_back(id, alloc);
		return alloc;
	}
}